"""
对训练集进行2D裁剪，每张ct抽取10张作为有标注数据，其他的作为无标注数据
并将标注和无标注的文件id用txt文件保存下来，2D-cut后的数据放在一个文件夹下
"""
import os
import random
import sys
sys.path.append(os.path.split(sys.path[0])[0])
import SimpleITK as sitk
from multiprocessing.dummy import Pool
from config import zunYiParameter as para

import numpy as np

# for file in tqdm(os.listdir(para.vol_path)[:para.train_test_split]):
def process(file):
    # 将CT和标签入读内存
    volume = sitk.ReadImage(os.path.join(para.nii_volume_path, file), sitk.sitkInt16)
    volume_array = sitk.GetArrayFromImage(volume)
    #print('处理前的volume shape:',volume_array.shape,end = '   ')
    seg = sitk.ReadImage(os.path.join(para.nii_seg_path, file.replace('volume', 'segmentation')), sitk.sitkUInt8)
    seg_array = sitk.GetArrayFromImage(seg)
    seg_file_name=file.replace('volume', 'segmentation')
    ct_file_name=file
    # seg_array=processLiverOrTumorMask(para.Liver_or_Tumor,seg_array)
    # 将灰度值在阈值之外的截断掉,MRI不用这个操作，ct才要
    # volume_array[volume_array > para.upper] = para.upper
    # volume_array[volume_array < para.lower] = para.lower

    all_choose_indexs=get_sampling(seg_array,filename = ct_file_name)#选取的所有下标
    all_indexs=[i for i in range(len(volume_array))]

    not_choose_indexs=list(set(all_indexs).difference(set(list(all_choose_indexs))))

    #保存图片，遍历有标注的下标
    choose_file_names=[]
    not_choose_file_names=[]
    for index_choose in all_choose_indexs:
        choose_ct_name=os.path.join("volume",str(index_choose)+"_"+ct_file_name)
        choose_seg_name=os.path.join("segmentation", str(index_choose) + "_" + seg_file_name)

        choose_file_names.append(choose_ct_name+" "+choose_seg_name)
        #保存
        ct_path = os.path.join(para.cut2d_save_path,choose_ct_name)
        seg_path = os.path.join(para.cut2d_save_path,choose_seg_name)

        sitk.WriteImage(sitk.GetImageFromArray(volume_array[index_choose]),ct_path)
        sitk.WriteImage(sitk.GetImageFromArray(seg_array[index_choose]), seg_path)
        # imageio.imwrite(ct_path, volume_array[index_choose])
        # imageio.imwrite(seg_path, seg_array[index_choose])
    for index_not_choose in not_choose_indexs:
        choose_ct_name = os.path.join("volume",str(index_not_choose)+"_"+ct_file_name)
        choose_seg_name = os.path.join("segmentation", str(index_not_choose) + "_" + seg_file_name)
        not_choose_file_names.append(choose_ct_name + " " + choose_seg_name)
        # 保存
        ct_path = os.path.join(para.cut2d_save_path, choose_ct_name)
        seg_path = os.path.join(para.cut2d_save_path, choose_seg_name)
        sitk.WriteImage(sitk.GetImageFromArray(volume_array[index_not_choose]), ct_path)
        sitk.WriteImage(sitk.GetImageFromArray(seg_array[index_not_choose]), seg_path)
    # [[选择作为训练的切片名称],[没有选择作为训练的切片名称]]
    return [choose_file_names,not_choose_file_names]

def get_sampling(seg_array,filename):
    '''
    遍历每一张，找出有肿瘤和无肿瘤的
    找到所有包含肿瘤的样本和少量背景切片。
    return:[],选择的切片下标列表
    '''
    no_tumor_num =999

    tumor_indexs_choose = list(np.where(np.any(seg_array == 1, axis = (1, 2)))[0])  # 肿瘤区域

    no_tumor_indexs = list(np.where(np.all(seg_array == 0, axis = (1, 2)))[0])  # 无肝脏，无肿瘤

    no_tumor_indexs_choose = random.sample(no_tumor_indexs,
                                           min(len(no_tumor_indexs),no_tumor_num))
    # 抽样
    print(filename+" 抽样：肿瘤，无肿瘤: ", len(tumor_indexs_choose), len(no_tumor_indexs_choose))
    return tumor_indexs_choose+ no_tumor_indexs_choose

def main():
    if not os.path.exists(para.cut2d_save_path_vol):
        os.makedirs(para.cut2d_save_path_vol)
    if not os.path.exists(para.cut2d_save_path_seg):
        os.makedirs(para.cut2d_save_path_seg)
    if not os.path.exists(os.path.dirname(para.train2d_choose_id_path)):
        os.makedirs(os.path.dirname(para.train2d_choose_id_path))
    if not os.path.exists(os.path.dirname(para.train2d_notchoose_id_path)):
        os.makedirs(os.path.dirname(para.train2d_notchoose_id_path))

    # 需要训练集的file_list
    with open(para.train_nii_id_path, 'r') as f:
        labeled_ids = f.read().splitlines()
    file_list = [os.path.basename(file.split()[0]) for file in labeled_ids]
    print("process nii file number:",len(file_list))
    # 多线程
    pool = Pool(4)
    result = pool.map(process, file_list)
    pool.close()
    pool.join()

    choose_file_names = [l[0] for l in result]#选包含肿瘤的，以及少量不包含肿瘤的
    not_choose_file_names = [l[1] for l in result]#大部分不包含肿瘤的背景区域

    choose_file_names_flatten = []
    not_choose_file_names_flatten = []
    for i in range(len(choose_file_names)):
        choose_file_names_flatten += choose_file_names[i]
        not_choose_file_names_flatten += not_choose_file_names[i]
    choose_file_names_flatten = np.array(choose_file_names_flatten)
    not_choose_file_names_flatten = np.array(not_choose_file_names_flatten)
    print("all choose", choose_file_names_flatten.shape)
    print("all not choose", not_choose_file_names_flatten.shape)
    np.savetxt(para.train2d_choose_id_path, choose_file_names_flatten, fmt = "%s")
    np.savetxt(para.train2d_notchoose_id_path, not_choose_file_names_flatten, fmt = "%s")

    print("save to txt file successful!")
if __name__ == '__main__':
    #
    # if os.path.exists(para.cut_save_path):
    #     shutil.rmtree(para.cut_save_path)
    '''
    抽取的切片的结果：
    G008.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  11 5
    G004.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  12 5
    G020.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  10 5
    G015.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  24 5
    G037.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  10 5
    G006.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  17 5
    G013.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  16 5
    G010.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  10 5
    G025.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  18 5
    G003.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  13 5
    G019.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  12 5
    G034.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  7 5
    G009.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  16 5
    G040.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  12 5
    G024.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  13 5
    G036.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  7 5
    G039.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  15 5
    G038.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  14 5
    G027.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  34 5
    G028.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  13 5
    G001.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  8 5
    G033.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  6 5
    G011.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  12 5
    G029.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  14 5
    G021.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  33 5
    G014.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  9 5
    G026.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  17 5
    G016.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  8 5
    G035.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  18 5
    G030.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  13 5
    G018.nii.gz 抽样：肿瘤，肝脏+无肿瘤，无肝脏无肿瘤:  12 5
    all choose (589,)
    all not choose (2491,)
    '''
    main()

